from pyspark.sql.functions import *
from pyspark.sql.session import SparkSession
from pyspark.sql.types import *
from pyspark.sql.functions import *
import jieba

# 创建spark sql执行环境
# enableHiveSupport:开启hive的元数据支持，可以读取到hive中的表
spark = SparkSession \
    .builder \
    .appName("rdd") \
    .config("spark.sql.shuffle.partitions", 1) \
    .enableHiveSupport() \
    .getOrCreate()

# 读取数据
comments_df = spark.read \
    .format("csv") \
    .schema("id STRING,comment_time STRING,comment STRING,a STRING ,b STRING") \
    .load("/data/comment.txt") \
    .where(col("comment").isNotNull())


# 自定义函数实现分词
# 1 定义方法
# 方法外属于Driver端
def cut_fun(comment):
    # 使用结巴分词器
    # 方法内属于Executor端
    return jieba.lcut(comment)


# 2 注册成spark函数
cut = udf(cut_fun, returnType=ArrayType(StringType()))

# 再DSL中使用自定义函数
words_df = comments_df \
    .select(cut("comment").alias("words"))

words_df.printSchema()

words_df \
    .select(explode("words").alias("word")) \
    .groupBy("word") \
    .agg(count("*").alias("num")) \
    .show()

# 再sql中使用自定义函数
# 注册表
comments_df.createOrReplaceTempView("comments")
# 注册sql函数
spark.udf.register("cut", cut_fun, ArrayType(StringType()))

spark.sql("""
select word,count(1) as num 
from comments
lateral view explode(cut(comment)) T as word
group by word
order by num desc
""").show()


# 代码上传到服务器运行
# 1  hdfs dfs -put comment.txt /data
# 2  pip install jieba  每台服务器都需要安装
# 3  spark-submit --master yarn  demo14_udf.py

